import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1=nn.Conv2d(1, 8, 5)
        self.pool1=nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(8, 16, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        # self.flatten=nn.Flatten()
        self.fc=nn.Linear(4*4*16, 10)
    def forward(self, x):
        x=self.pool1(F.relu(self.conv1(x)))
        x=self.pool1(F.relu(self.conv2(x)))
        x=x.view(x.shape[0], -1)
        x=self.fc(F.relu(x))
        return x